#include <iostream>

using namespace std;

#include <IL/il.h>

#define BLOCK   16
#define STEP     4
#define STREAMS  8 //nombre de streams à lancer

struct complex {

  float r;
  float i;

  __device__ complex(float r, float i) : r(r), i(i) {}

  __device__ float magnitude() {return r*r + i*i;}

  __device__ complex operator*(const complex& c) {
    return complex(r * c.r - i * c.i, i * c.r + r * c.i);
  }

  __device__ complex operator+(const complex& c) {
    return complex(r + c.r, i + c.i);
  }  
  
  __device__ void iter(const complex &c){
    float aux = r;
    r = r*r - i*i + c.r;
    aux = i*aux;
    i = aux+aux     + c.i;
  }
};


__device__ int julia(int x, int y, int size) {

  const float scale = 1.5;

  float jx = scale * (float)(size/2 - x)/(size/2);
  float jy = scale * (float)(size/2 - y)/(size/2);

  //const complex c(0.285, 0.013);
  const complex c(-0.8, 0.156);

  complex a(jx, jy);

  bool sortir(false);

  for(unsigned int i(0); i < 25 && !sortir; i++) {
    a.iter(c);
    sortir |= (a.magnitude() > 1000);
    a.iter(c);
    sortir |= (a.magnitude() > 1000);
    a.iter(c);
    sortir |= (a.magnitude() > 1000);
    a.iter(c);
    sortir |= (a.magnitude() > 1000);
    a.iter(c);
    sortir |= (a.magnitude() > 1000);
    a.iter(c);
    sortir |= (a.magnitude() > 1000);
    a.iter(c);
    sortir |= (a.magnitude() > 1000);
    a.iter(c);
    sortir |= (a.magnitude() > 1000);

    if(sortir) return 0;
  }

  return 1;

}

//on ajoute h, un paramètre servant à définir la position du stream dans l'image
//sur l'axe des y.
__global__ void kernel(unsigned char* image, unsigned int size, int h) {
  //ici rien ne change
  unsigned int x = blockIdx.x*blockDim.x+threadIdx.x;
  unsigned int y = (blockIdx.y*blockDim.y+threadIdx.y)*STEP;

  __shared__ unsigned char pixels[BLOCK*STEP][BLOCK][4];

#pragma unroll
  for(unsigned int k = 0; k < STEP; ++k){
    pixels[threadIdx.y*STEP+k][threadIdx.x][0] = 0;
    pixels[threadIdx.y*STEP+k][threadIdx.x][1] = 0;
    //en revanche ici il faut dire à Julia notre vrai position
    pixels[threadIdx.y*STEP+k][threadIdx.x][2] = 255*julia(x,y+k+h,size);
    pixels[threadIdx.y*STEP+k][threadIdx.x][3] = 255;
  }

#pragma unroll
  for(unsigned int k = 0; k < STEP; ++k){
    unsigned int offset = ((y+k)*size+x);
    reinterpret_cast<int*>(image)[offset]
      = reinterpret_cast<int*>(pixels[threadIdx.y*STEP+k])[threadIdx.x];
    
  }
}

__global__ void blanc(unsigned char* image, unsigned int size) {
  unsigned int x = blockIdx.x*blockDim.x+threadIdx.x;
  unsigned int y = (blockIdx.y*blockDim.y+threadIdx.y)*STEP;

  for(unsigned int k = 0; k < STEP; k++){
    unsigned int offset = ((y+k)*size+x)*4;
    image[offset  ] = 255;
    image[offset+1] = 255;
    image[offset+2] = 255;
    image[offset+3] = 255;
  }
}


int main(int argc, char* argv[]) {
  cudaEvent_t start, stop0, stop1;
  float time;

  cudaEventCreate(&start);
  cudaEventCreate(&stop0);
  cudaEventCreate(&stop1);
  
  //la taille doit être maintenant, en plus, multiple de STREAMS
  unsigned int size = static_cast<unsigned int>(atoi(argv[1]));
  size = ((size/(BLOCK*STEP*STREAMS))+(size%(BLOCK*STEP*STREAMS) != 0))
    *BLOCK*STEP*STREAMS;
  
  //la hauteur d'un bloc est la hauteur totale divisée par STREAMS
  unsigned int hblock = size/STREAMS;

  //la grille est maintenant dépendante de hblock en hauteur
  dim3 grid(size/BLOCK, hblock/(BLOCK*STEP));
  dim3 block(BLOCK, BLOCK);
  
  //on alloue l'image en version cudaMallocHost
  unsigned char* image;
  cudaMallocHost(&image, size*size*4);
 
  //on alloue un buffer par stream
  unsigned char* dev_image[STREAMS];
  for(uint i(0); i < STREAMS; ++i){
    cudaMalloc(&dev_image[i], hblock*size*4);
    //blanc<<<grid,block>>>(dev_image[i], size);
  }
  
  //====================================================================
  // On crée les streams
  cudaStream_t streams[STREAMS];
  for(uint i(0); i < STREAMS; ++i)
    cudaStreamCreate(&streams[i]);
  //*/
  cudaEventRecord(start, 0);
  
  //on lance les kernels sur les streams
  for(uint i(0); i < STREAMS; ++i)
    kernel<<<grid,block,0,streams[i]>>>(dev_image[i], size, i*hblock);
  
  //les copies
  for(uint i(0); i < STREAMS; ++i)
    cudaMemcpyAsync(&image[i*hblock*size*4], 
		    dev_image[i], hblock*size*4, 
		    cudaMemcpyDeviceToHost, streams[i]);
  

  cudaEventRecord(stop1, 0);

  cudaEventSynchronize(stop1);

  cudaEventElapsedTime(&time, start, stop1);
  cout << "time total=" << time << endl;

  //on détruit les streams
  for(uint i(0); i < STREAMS; ++i)
    cudaStreamDestroy(streams[i]);
    

  cudaEventDestroy(start);
  cudaEventDestroy(stop0);
  cudaEventDestroy(stop1);

  unsigned int handle;
  
  //et les buffer GPU
  for(uint i(0); i < STREAMS; ++i)
    cudaFree(dev_image[i]);

  ilInit();
  ilGenImages(1, &handle);
  ilBindImage(handle);
  ilTexImage(size, size, 1, 4, IL_RGBA, IL_UNSIGNED_BYTE, image);
  ilSetData(image);
  ilEnable(IL_FILE_OVERWRITE);
  ilSaveImage("out-cuda.png");
  ilDeleteImages(1, &handle);

  //puis l'image
  cudaFreeHost(image);

  return 0;

}
